#!/usr/bin/env python3
"""
心跳上报客户端Demo
模拟NodeManager向Controller发送心跳消息，报告其管理的endpoint状态
"""

import requests
import json
import time
import threading
import random
from typing import Dict, List
from motor.resources.endpoint import EndpointStatus


class HeartbeatClient:
    """心跳客户端类"""
    
    def __init__(self, controller_url: str = "http://localhost:57675", 
                 instance_id: int = 1, job_name: str = "test_job", 
                 pod_ip: str = "127.0.0.1"):
        self.controller_url = controller_url
        self.instance_id = instance_id
        self.job_name = job_name
        self.pod_ip = pod_ip
        self.heartbeat_url = f"{controller_url}/controller/heartbeat"
        
        # 模拟的endpoint状态
        self.endpoints_status = {
            0: EndpointStatus.NORMAL,
            1: EndpointStatus.NORMAL,
        }
        
        self.running = False
        self.heartbeat_thread = None
        
    def create_heartbeat_message(self) -> Dict:
        """创建心跳消息"""
        return {
            "job_name": self.job_name,
            "ins_id": self.instance_id,
            "ip": self.pod_ip,
            "status": {k: v.value for k, v in self.endpoints_status.items()}
        }
    
    def send_heartbeat(self) -> bool:
        """发送心跳消息"""
        try:
            heartbeat_msg = self.create_heartbeat_message()
            
            print(f"发送心跳消息: {json.dumps(heartbeat_msg, indent=2, ensure_ascii=False)}")
            
            response = requests.post(
                self.heartbeat_url, 
                json=heartbeat_msg, 
                timeout=10
            )
            
            if response.status_code == 200:
                result = response.json()
                print(f"心跳发送成功! 响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
                return True
            else:
                print(f"心跳发送失败! 状态码: {response.status_code}")
                print(f"错误信息: {response.text}")
                return False
                
        except requests.exceptions.ConnectionError:
            print("连接失败! 请确保控制器服务正在运行")
            return False
        except requests.exceptions.Timeout:
            print("心跳请求超时!")
            return False
        except Exception as e:
            print(f"发送心跳时发生错误: {e}")
            return False
    
    def simulate_endpoint_status_change(self):
        """模拟endpoint状态变化"""
        if random.random() < 0.1:  # 10%概率改变状态
            endpoint_id = random.choice(list(self.endpoints_status.keys()))
            current_status = self.endpoints_status[endpoint_id]
            
            # 随机选择新状态
            if current_status == EndpointStatus.NORMAL:
                new_status = random.choice([EndpointStatus.ABNORMAL, EndpointStatus.PAUSED])
            elif current_status == EndpointStatus.ABNORMAL:
                new_status = random.choice([EndpointStatus.NORMAL, EndpointStatus.PAUSED])
            elif current_status == EndpointStatus.PAUSED:
                new_status = random.choice([EndpointStatus.NORMAL, EndpointStatus.ABNORMAL])
            else:
                new_status = EndpointStatus.NORMAL
                
            self.endpoints_status[endpoint_id] = new_status
            print(f"Endpoint {endpoint_id} 状态从 {current_status.value} 变为 {new_status.value}")
    
    def heartbeat_loop(self, interval: float = 2.0):
        """心跳循环"""
        print(f"开始发送心跳，间隔: {interval}秒")
        
        while self.running:
            # 模拟endpoint状态变化
            # self.simulate_endpoint_status_change()
            
            # 发送心跳
            success = self.send_heartbeat()
            
            if not success:
                print("心跳发送失败，等待重试...")
                time.sleep(1)
                continue
            
            # 等待下次心跳
            time.sleep(interval)
    
    def start(self, interval: float = 2.0):
        """启动心跳客户端"""
        if self.running:
            print("心跳客户端已经在运行中")
            return
            
        self.running = True
        self.heartbeat_thread = threading.Thread(
            target=self.heartbeat_loop, 
            args=(interval,),
            daemon=True,
            name="HeartbeatClient"
        )
        self.heartbeat_thread.start()
        print(f"心跳客户端已启动，实例ID: {self.instance_id}, Job名称: {self.job_name}")
    
    def stop(self):
        """停止心跳客户端"""
        if not self.running:
            print("心跳客户端未运行")
            return
            
        self.running = False
        if self.heartbeat_thread:
            self.heartbeat_thread.join(timeout=5)
        print("心跳客户端已停止")
    
    def set_endpoint_status(self, endpoint_id: int, status: EndpointStatus):
        """手动设置endpoint状态"""
        if endpoint_id in self.endpoints_status:
            old_status = self.endpoints_status[endpoint_id]
            self.endpoints_status[endpoint_id] = status
            print(f"手动设置 Endpoint {endpoint_id} 状态从 {old_status.value} 变为 {status.value}")
        else:
            print(f"Endpoint {endpoint_id} 不存在")


def single_heartbeat_example():
    """单次心跳发送示例"""
    print("=== 单次心跳发送示例 ===")
    
    client = HeartbeatClient(
        instance_id=1,
        job_name="test_job",
        pod_ip="127.0.0.1"
    )
    
    # 发送一次心跳
    client.send_heartbeat()


def continuous_heartbeat_example():
    """持续心跳发送示例"""
    print("=== 持续心跳发送示例 ===")
    
    client = HeartbeatClient(
        instance_id=1,
        job_name="test_job",
        pod_ip="127.0.0.1"
    )
    
    try:
        # 启动心跳客户端
        client.start(interval=3.0)  # 每3秒发送一次心跳
        
        # 运行30秒
        print("心跳客户端将运行30秒...")
        time.sleep(30)
        
        # 模拟手动改变endpoint状态
        print("\n模拟endpoint状态变化...")
        client.set_endpoint_status(0, EndpointStatus.ABNORMAL)
        time.sleep(5)
        
        client.set_endpoint_status(0, EndpointStatus.NORMAL)
        time.sleep(5)
        
        client.set_endpoint_status(1, EndpointStatus.PAUSED)
        time.sleep(5)
        
    except KeyboardInterrupt:
        print("\n收到中断信号，正在停止...")
    finally:
        client.stop()


def multi_instance_heartbeat_example():
    """多实例心跳发送示例"""
    print("=== 多实例心跳发送示例 ===")
    
    # 创建多个心跳客户端
    clients = []
    for i in range(3):
        client = HeartbeatClient(
            instance_id=i + 1,
            job_name=f"test_job_{i + 1}",
            pod_ip=f"127.0.0.{i + 1}"
        )
        clients.append(client)
        client.start(interval=2.0 + i)  # 不同的心跳间隔
    
    try:
        print("多实例心跳客户端已启动，将运行20秒...")
        time.sleep(20)
        
    except KeyboardInterrupt:
        print("\n收到中断信号，正在停止所有客户端...")
    finally:
        for client in clients:
            client.stop()


def interactive_heartbeat_example():
    """交互式心跳示例"""
    print("=== 交互式心跳示例 ===")
    print("输入命令来控制心跳客户端:")
    print("  start - 启动心跳")
    print("  stop - 停止心跳")
    print("  status <endpoint_id> <status> - 设置endpoint状态")
    print("  quit - 退出")
    
    client = HeartbeatClient(
        instance_id=1,
        job_name="interactive_job",
        pod_ip="127.0.0.1"
    )
    
    try:
        while True:
            command = input("\n请输入命令: ").strip().split()
            
            if not command:
                continue
                
            if command[0] == "start":
                client.start(interval=2.0)
            elif command[0] == "stop":
                client.stop()
            elif command[0] == "status" and len(command) >= 3:
                try:
                    endpoint_id = int(command[1])
                    status_name = command[2].upper()
                    if hasattr(EndpointStatus, status_name):
                        status = getattr(EndpointStatus, status_name)
                        client.set_endpoint_status(endpoint_id, status)
                    else:
                        print(f"无效的状态: {status_name}")
                except ValueError:
                    print("无效的endpoint ID")
            elif command[0] == "quit":
                break
            else:
                print("未知命令")
                
    except KeyboardInterrupt:
        print("\n收到中断信号...")
    finally:
        client.stop()


if __name__ == "__main__":
    print("心跳客户端Demo")
    print("请选择运行模式:")
    print("1. 单次心跳发送")
    print("2. 持续心跳发送")
    print("3. 多实例心跳发送")
    print("4. 交互式心跳")
    
    try:
        choice = input("请输入选择 (1-4): ").strip()
        
        if choice == "1":
            single_heartbeat_example()
        elif choice == "2":
            continuous_heartbeat_example()
        elif choice == "3":
            multi_instance_heartbeat_example()
        elif choice == "4":
            interactive_heartbeat_example()
        else:
            print("无效选择，运行单次心跳示例")
            single_heartbeat_example()
            
    except KeyboardInterrupt:
        print("\n程序被中断")
    except Exception as e:
        print(f"程序发生错误: {e}")
